package org.example.erp.utils;

import org.example.erp.myInterface.NodeMapper;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;

/*
 * @Date 2024-10-20
 * @Author 28254
 * @ClassName TreeBuild
 * 用于构建树形数据的工具类
 */

public class TreeBuild<T> {
    /**
     * 构建树形数据方法
     *
     * @param nodes
     * @param nodeMapper
     * @return
     */
    public List<T> buildTree(List<T> nodes, NodeMapper<T> nodeMapper) {
        Map<Integer, T> nodeMap = new HashMap<>();
        List<T> rootNodes = new ArrayList<>();
        // 将所有节点放入map中
        for (T node : nodes) {
            nodeMap.put(nodeMapper.getId(node), node);
        }
        // 构建树形菜单
        for (T node : nodes) {
            // 如果是根节点 或者 找不到父节点，则认为是根节点
            if (nodeMapper.getParentId(node) == null || !nodeMap.containsKey(nodeMapper.getParentId(node))) {
                rootNodes.add(node);
            } else {
                // 否则，找到其父节点，并将其添加到父节点的子节点列表中
                T parentNode = nodeMap.get(nodeMapper.getParentId(node));
                List<T> children = nodeMapper.getChildren(parentNode);
                //如果没有则创建一个新的子节点
                if (children == null || children.isEmpty()) {
                    children = new ArrayList<>();
                    nodeMapper.setChildren(parentNode, children);
                }
                children.add(node);
            }
        }
        return rootNodes;
    }

    // 通用实现nodeMapper
    public static <T> NodeMapper<T> createNodeMapper(
            Function<T, Integer> idGetter,
            Function<T, Integer> parentIdGetter,
            Function<T, List<T>> childrenGetter,
            BiConsumer<T, List<T>> childrenSetter
    ) {
        return new NodeMapper<>() {
            @Override
            public Integer getId(T node) {
                return idGetter.apply(node);
            }

            @Override
            public Integer getParentId(T node) {
                return parentIdGetter.apply(node);
            }

            @Override
            public List<T> getChildren(T node) {
                return childrenGetter.apply(node);
            }

            @Override
            public void setChildren(T node, List<T> children) {
                childrenSetter.accept(node, children);
            }
        };
    }
}
